import numpy as np
import torch.nn as nn
import torch
from torch.autograd import Variable
from matplotlib import pyplot as plot


class AutoEncoderNetwork(nn.Module):

    def __init__(self):
        super(AutoEncoderNetwork, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(18, 12, bias=True),
            nn.Tanh(),
            nn.Linear(12, 9, bias=True),
            nn.Tanh(),
            nn.Linear(9, 5, bias=True),
            nn.Tanh(),
            nn.Linear(5, 3, bias=True),
            nn.Tanh()
        )
        self.decoder = nn.Sequential(
            nn.Linear(3, 5, bias=True),
            nn.Tanh(),
            nn.Linear(5, 9, bias=True),
            nn.Tanh(),
            nn.Linear(9, 13, bias=True),
            nn.Tanh(),
            nn.Linear(13, 18, bias=True),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.decoder(self.encoder(x))


def main():
    ae = AutoEncoderNetwork().cuda()

    batch_size = 1
    sample_size = 5
    vec_l = 18

    train_data_numpy = np.random.rand(batch_size * sample_size * vec_l).reshape([-1, vec_l])
    train_tensor = torch.DoubleTensor(train_data_numpy)
    x = Variable(torch.FloatTensor(train_tensor.numpy())).cuda()
    optimizer = torch.optim.Adam(ae.parameters(), 1e-2)
    loss_fun = nn.MSELoss()

    step_info = []
    loss_info = []

    for i in range(100000):
        y = ae.forward(x)
        loss = loss_fun(y, x)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if i % 1000 == 0:
            step_info.append(i)
            loss_info.append(loss.cpu().data.numpy())
            plot.plot(step_info, loss_info)
            print("current loss: %f" % loss.cpu().data)
            print("current index :%d" % i)
            print("raw:" + str(x.cpu().data.numpy()))
            print("encode:" + str(ae.encoder(x).cpu().data.numpy()))
            print("decode:" + str(ae.decoder(ae.encoder(x)).cpu().data.numpy()))
            print("sub:" + str((torch.max(torch.abs(ae.decoder(ae.encoder(x)) - x))).cpu().data))
            plot.show()


if __name__ == "__main__":
    main()
